{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bda38497",
   "metadata": {},
   "source": [
    "## Hurdle and truncated count models\n",
    "\n",
    "Author: Josef Perktold\n",
    "\n",
    "Statsmodels has now hurdle and truncated count models, added in version 0.14.\n",
    "\n",
    "A hurdle model is composed of a model for zeros and a model for the distribution for counts larger than zero. The zero model is a binary model for a count of zero versus larger than zero. The count model for nonzero counts is a zero truncated count model.\n",
    "\n",
    "Statsmodels currently supports hurdle models with Poisson and Negative Binomial distributions as zero model and as count model. Binary models like Logit, Probit or GLM-Binomial are not yet supported as zero model.\n",
    "The advantage of Poisson-Poisson hurdle is that the standard Poisson model is a special case with equal parameters in both models. This provides a simple Wald test for the hurdle model against the Poisson model.\n",
    "\n",
    "The implemented binary model is a censored model where observations are right censored at one. That means that only 0 or 1 counts are observed.\n",
    "\n",
    "The hurdle model can be estimated by separately estimating the zero model and the count model for the zero truncated data assuming that observations are independently distributed (no correlation across observations). The resulting covariance matrix of the parameter estimates is block diagonal with diagonal blocks given by the submodels.\n",
    "Joint estimation is not yet implemented.\n",
    "\n",
    "The censored and truncated count models were developed mainly to support the hurdle model. However, the left truncated count models have other applications than supporting the hurdle models. The right censored models are not of separate interest because they only support binary observations that can be modeled by GLM-Binomial, Logit or Probit.\n",
    "\n",
    "For the hurdle model there is a single class `HurdleCountModel`, that includes the distributions of the submodels as option. \n",
    "Classes for truncated models are currently `TruncatedLFPoisson` and `TruncatedLFNegativeBinomialP`, where \"LF\" stands for left truncation at a fixed, observation independent truncation point. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29e6105a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eed890e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import statsmodels.discrete.truncated_model as smtc\n",
    "from statsmodels.discrete.count_model import (\n",
    "    ZeroInflatedGeneralizedPoisson,\n",
    "    ZeroInflatedNegativeBinomialP,\n",
    "    ZeroInflatedPoisson,\n",
    ")\n",
    "from statsmodels.discrete.discrete_model import (\n",
    "    GeneralizedPoisson,\n",
    "    NegativeBinomial,\n",
    "    NegativeBinomialP,\n",
    "    Poisson,\n",
    ")\n",
    "from statsmodels.discrete.truncated_model import (\n",
    "    HurdleCountModel,\n",
    "    TruncatedLFNegativeBinomialP,\n",
    "    TruncatedLFPoisson,\n",
    "    _RCensoredPoisson,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85da298e",
   "metadata": {},
   "source": [
    "## Simulating a hurdle model\n",
    "\n",
    "We are simulating a Poisson-Poisson hurdle model explicitly because there are not yet any distribution helper functions for it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99c1d162",
   "metadata": {},
   "outputs": [],
   "source": [
    "np.random.seed(987456348)\n",
    "# large sample to get strong results\n",
    "nobs = 5000\n",
    "x = np.column_stack((np.ones(nobs), np.linspace(0, 1, nobs)))\n",
    "\n",
    "mu0 = np.exp(0.5 * 2 * x.sum(1))\n",
    "y = np.random.poisson(mu0, size=nobs)\n",
    "print(np.bincount(y))\n",
    "y_ = y\n",
    "indices = np.arange(len(y))\n",
    "mask = mask0 = y > 0\n",
    "for _ in range(10):\n",
    "\n",
    "    print(mask.sum())\n",
    "    indices = mask  # indices[mask]\n",
    "    if not np.any(mask):\n",
    "        break\n",
    "    mu_ = np.exp(0.5 * x[indices].sum(1))\n",
    "    y[indices] = y_ = np.random.poisson(mu_, size=len(mu_))\n",
    "    np.place(y, mask, y_)\n",
    "    mask = np.logical_and(mask0, y == 0)\n",
    "\n",
    "np.bincount(y)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6435d35",
   "metadata": {},
   "source": [
    "## Estimating misspecified Poisson Model\n",
    "\n",
    "The data that we generated has zero deflation, this is, we observe fewer zeros than what we would expect in a Poisson model.\n",
    "\n",
    "After fitting the model, we can use the plot function in the poisson diagnostic class to compare the expected predictive distribution and the realized frequencies. The shows that the Poisson model overestimates the number of zeros and underestimates counts of one and two."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7fdd59b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "mod_p = Poisson(y, x)\n",
    "res_p = mod_p.fit()\n",
    "print(res_p.summary())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91a5c9cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "dia_p = res_p.get_diagnostic()\n",
    "dia_p.plot_probs();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a237e21f",
   "metadata": {},
   "source": [
    "## Estimating the Hurdle Model\n",
    "\n",
    "Next, we estimate the correctly specified Poisson-Poisson hurdle model.\n",
    "\n",
    "Signature and options for the HurdleCountModel shows that poisson-poisson is the default, so we do not need to specify any options when creating this model.\n",
    "\n",
    "`HurdleCountModel(endog, exog, offset=None, dist='poisson', zerodist='poisson', \n",
    "                  p=2, pzero=2, exposure=None, missing='none', **kwargs)`\n",
    "                  \n",
    "The results class of the HurdleCountModel has a `get_diagnostic` method. However, only part of the diagnostic methods are currently available. The plot of the predictive distribution shows very high agreement with the data.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70065ce8",
   "metadata": {},
   "outputs": [],
   "source": [
    "mod_h = HurdleCountModel(y, x)\n",
    "res_h = mod_h.fit(disp=False)\n",
    "print(res_h.summary())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c3aadbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "dia_h = res_h.get_diagnostic()\n",
    "dia_h.plot_probs();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0827fde3",
   "metadata": {},
   "source": [
    "We can use the Wald test to test whether the parameters of the zero model are the same as the parameters of the zero-truncated count model. The p-value is very small and correctly rejects that the model is just Poisson. We are using a large sample size, so the power of the test will be large in this case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c000b4cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_h.wald_test(\"zm_const = const, zm_x1 = x1\", scalar=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "100ba8a8",
   "metadata": {},
   "source": [
    "## Prediction\n",
    "\n",
    "The hurdle model can be used for prediction for statistics of the overall model and of the two submodels. The statistics that should be predicted is specified using the `which` keyword.\n",
    "\n",
    "The following is taken from the docstring for predict and lists available the options.\n",
    "\n",
    "        which : str (optional)\n",
    "            Statitistic to predict. Default is 'mean'.\n",
    "\n",
    "            - 'mean' : the conditional expectation of endog E(y | x)\n",
    "            - 'mean-main' : mean parameter of truncated count model.\n",
    "              Note, this is not the mean of the truncated distribution.\n",
    "            - 'linear' : the linear predictor of the truncated count model.\n",
    "            - 'var' : returns the estimated variance of endog implied by the\n",
    "              model.\n",
    "            - 'prob-main' : probability of selecting the main model which is\n",
    "              the probability of observing a nonzero count P(y > 0 | x).\n",
    "            - 'prob-zero' : probability of observing a zero count. P(y=0 | x).\n",
    "              This is equal to is ``1 - prob-main``\n",
    "            - 'prob-trunc' : probability of truncation of the truncated count\n",
    "              model. This is the probability of observing a zero count implied\n",
    "              by the truncation model.\n",
    "            - 'mean-nonzero' : expected value conditional on having observation\n",
    "              larger than zero, E(y | X, y>0)\n",
    "            - 'prob' : probabilities of each count from 0 to max(endog), or\n",
    "              for y_values if those are provided. This is a multivariate\n",
    "              return (2-dim when predicting for several observations).\n",
    "              \n",
    "These options are available in the `predict` and the `get_prediction` methods of the results class.\n",
    "\n",
    "For the following example, we create a set of explanatory variables that are taken from the original data at equal spaced intervals. Then we can predict the available statistics conditional on these explanatory variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9757da68",
   "metadata": {},
   "outputs": [],
   "source": [
    "which_options = [\n",
    "    \"mean\",\n",
    "    \"mean-main\",\n",
    "    \"linear\",\n",
    "    \"mean-nonzero\",\n",
    "    \"prob-zero\",\n",
    "    \"prob-main\",\n",
    "    \"prob-trunc\",\n",
    "    \"var\",\n",
    "    \"prob\",\n",
    "]\n",
    "ex = x[slice(None, None, nobs // 5), :]\n",
    "ex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3614ee25",
   "metadata": {},
   "outputs": [],
   "source": [
    "for w in which_options:\n",
    "    print(w)\n",
    "    pred = res_h.predict(ex, which=w)\n",
    "    print(\"    \", pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbc9ebed",
   "metadata": {},
   "outputs": [],
   "source": [
    "for w in which_options[:-1]:\n",
    "    print(w)\n",
    "    pred = res_h.get_prediction(ex, which=w)\n",
    "    print(\"    \", pred.predicted)\n",
    "    print(\"  se\", pred.se)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e356eb1b",
   "metadata": {},
   "source": [
    "The option `which=\"prob\"` returns an array of predicted probabilities for each row of the predict `exog`.\n",
    "We are often interested in the mean probabilities averaged over all exog. The prediction methods have an option `average=True` to compute the average of the predicted values across observations and the corresponding standard errors and confidence intervals for those averaged predictions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d90afc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred = res_h.get_prediction(ex, which=\"prob\", average=True)\n",
    "print(\"    \", pred.predicted)\n",
    "print(\"  se\", pred.se)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09602e97",
   "metadata": {},
   "source": [
    "We use the panda DataFrame to get a display that is easier to read. The \"predicted\" column shows the probability mass function for the predicted distribution of response values averaged of our 5 grid points of exog. The probabilities do not add up to one because counts larger than those observed have positive probability and are missing in the table, although in this example that probability is small."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77c7eff9",
   "metadata": {},
   "outputs": [],
   "source": [
    "dfp_h = pred.summary_frame()\n",
    "dfp_h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "185a5042",
   "metadata": {},
   "outputs": [],
   "source": [
    "prob_larger9 = pred.predicted.sum()\n",
    "prob_larger9, 1 - prob_larger9"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a75795cd",
   "metadata": {},
   "source": [
    "`get_prediction` returns in this case an instance of the base `PredictionResultsDelta` class.\n",
    "\n",
    "Inferential statistics like standard errors, p-values and confidence interval for nonlinear functions that depend on several distribution parameters are computed using the delta method. Inference for predictions is based on the normal distribution. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8124e02",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c1d314d",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred.dist, pred.dist_args"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b032002",
   "metadata": {},
   "source": [
    "We can compare the distribution predicted by the hurdle model with the one predicted by the Poisson model that we estimated earlier. The last column, \"diff\", shows that Poisson model overestimates the number of zeros by around 8% of observations and underestimates the counts of 1 and 2 by 7%, resp. 3.7% at the average over the `exog` grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd01ba63",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_p = res_p.get_prediction(ex, which=\"prob\", average=True)\n",
    "dfp_p = pred_p.summary_frame()\n",
    "dfp_h[\"poisson\"] = dfp_p[\"predicted\"]\n",
    "dfp_h[\"diff\"] = dfp_h[\"poisson\"] - dfp_h[\"predicted\"]\n",
    "dfp_h"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0da7fb38",
   "metadata": {},
   "source": [
    "## Other post-estimation\n",
    "\n",
    "The estimated hurdle model can be use for wald test of parameters and for prediction. Other maximum likelihood statistics such as loglikelihood value and information criteria are also available. \n",
    "\n",
    "However, some post-estimation methods that require helper functions that are not needed for estimation, parameter inference and prediction are not yet available. The main methods that are not supported yet are `score_test`, `get_distribution`, and `get_influence`. Diagnostic measures in `get_diagnostics` are only available for statistics that are based on prediction.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e821f6e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_h.llf, res_h.df_resid, res_h.aic, res_h.bic"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31d2ae62",
   "metadata": {},
   "source": [
    "Is there excess dispersion? We can use the pearson residuals to compute a pearson chi2 statistics which should be close to 1 if the model is correctly specified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f1a5124",
   "metadata": {},
   "outputs": [],
   "source": [
    "(res_h.resid_pearson**2).sum() / res_h.df_resid"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ff9df60",
   "metadata": {},
   "source": [
    "The diagnostic class also has the predictive distribution which is used in the diagnostic plots. No other statistics or tests are currently availalbe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50ba5545",
   "metadata": {},
   "outputs": [],
   "source": [
    "dia_h.probs_predicted.mean(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c913a1bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_h.resid[:10]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
